#include "process.h"

STATIC DWORD_PTR PsSystemPageTable[0x100];

HANDLE hProcesses[0x100];
DWORD dwProcessCount;

void PsInitializeSystemPageTable(void) {
  register DWORD i, j;
  register LPDWORD lpdwPage;
  dwProcessCount = 0;
  for (i = 0; i < 0x100; i++) {
    lpdwPage = (LPDWORD)(0x2000000 + (i << 12));
    PsSystemPageTable[i] = (DWORD)lpdwPage | PG_US_U | PG_RW_W | PG_P;
    for (j = 0; j < 0x400; j++) {
      lpdwPage[j] = ((i << 22) | (j << 12) | PG_US_U | PG_RW_W | PG_P);
    }
  }
}

HANDLE PsStartProcess(LPVOID lpExecutable) {
  HANDLE hProcess;
  LPPROCESS_CONTEXT lpPsContext;
  PIMAGE_DOS_HEADER pDosHeader;
  PIMAGE_NT_HEADERS32 pNtHeader;
  PIMAGE_SECTION_HEADER pSection;
  register DWORD i, j;
  register DWORD dwAddress;
  if (dwProcessCount >= 0x100) {
    return NULL;
  }
  pDosHeader = (PIMAGE_DOS_HEADER)lpExecutable;
  if (pDosHeader->e_magic != 0x5A4D) {
    return NULL;
  }
  pNtHeader =
      (PIMAGE_NT_HEADERS32)((DWORD_PTR)lpExecutable + pDosHeader->e_lfanew);
  if (pNtHeader->Signature != 0x00004550 ||
      pNtHeader->FileHeader.SizeOfOptionalHeader !=
          sizeof(IMAGE_OPTIONAL_HEADER32)) {
    return NULL;
  }
  hProcess = HeapAllocate(sizeof(struct PROCESS_CONTEXT));
  lpPsContext = (LPPROCESS_CONTEXT)hProcess;
  RtlZeroMemory(lpPsContext, sizeof(struct PROCESS_CONTEXT));
  lpPsContext->Signature = HANDLE_PROCESS;
  lpPsContext->cr3 = MmAllocatePage();
  PsInitializePageDirectory((LPDWORD)lpPsContext->cr3);
  pSection = (PIMAGE_SECTION_HEADER)((DWORD_PTR)pNtHeader +
                                     sizeof(IMAGE_NT_HEADERS32));
  for (i = 0; i < pNtHeader->FileHeader.NumberOfSections; i++) {
    for (j = 0; j < pSection->SizeOfRawData; j += 0x200) {
      dwAddress =
          pNtHeader->OptionalHeader.ImageBase + pSection->VirtualAddress + j;
      RtlCopyMemory(
          (LPBYTE)PsAllocatePage((LPDWORD)lpPsContext->cr3, dwAddress) +
              (j & 0xFFF),
          (LPBYTE)lpExecutable + pSection->PointerToRawData + j, 0x200);
    }
    ++pSection;
  }
  ThCreateProcessThread(
      hProcess, (LPVOID)(pNtHeader->OptionalHeader.ImageBase +
                         pNtHeader->OptionalHeader.AddressOfEntryPoint));
  hProcesses[dwProcessCount++] = hProcess;
  return hProcess;
}

void PsInitializePageDirectory(LPDWORD lpdwPageDir) {
  register DWORD i;
  RtlZeroMemory(lpdwPageDir, 0x1000);
  for (i = 0; i < 0x100; i++) {
    lpdwPageDir[i + 0x300] = PsSystemPageTable[i];
  }
}

LPDWORD PsAllocatePage(LPDWORD lpdwPageDir, DWORD_PTR dwAddress) {
  register LPDWORD lpdwPageTable;
  register LPDWORD lpdwPage;
  if (dwAddress > 0xc0000000) {
    return NULL;
  }
  lpdwPageDir += (dwAddress >> 22);
  if (!(*lpdwPageDir & PG_P)) {
    lpdwPage = (LPDWORD)MmAllocatePage();
    *lpdwPageDir = (DWORD)lpdwPage | PG_US_U | PG_RW_W | PG_P;
    RtlZeroMemory(lpdwPage, 0x1000);
  }
  lpdwPageTable =
      ((LPDWORD)(*lpdwPageDir & 0xFFFFF000)) + ((dwAddress >> 12) & 0x3FF);
  if (!(*lpdwPageTable & PG_P)) {
    lpdwPage = (LPDWORD)MmAllocatePage();
    *lpdwPageTable = (DWORD)lpdwPage | PG_US_U | PG_RW_W | PG_P;
    return lpdwPage;
  }
  return (LPDWORD)(*lpdwPageTable & 0xFFFFF000);
}

void PsFreePage(LPDWORD lpdwPageDir, DWORD_PTR dwAddress) {
  register LPDWORD lpdwPageTable;
  if (dwAddress > 0xc0000000) {
    return;
  }
  lpdwPageDir += (dwAddress >> 22);
  if (!(*lpdwPageDir & PG_P)) {
    return;
  }
  lpdwPageTable =
      ((LPDWORD)(*lpdwPageDir & 0xFFFFF000)) + ((dwAddress >> 12) & 0x3FF);
  if (*lpdwPageTable & PG_P) {
    MmFreePage(*lpdwPageTable & 0xFFFFF000);
    *lpdwPageTable = 0;
  }
}
